import aiohttp
import json
import time
import asyncio
from pydify import WorkflowClient
from pydify.workflow import WorkflowEvent
from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type

class WorkflowClientManager:
    def __init__(self, api_key: str, base_url: str, session_pool: aiohttp.ClientSession):
        self.client = WorkflowClient(api_key=api_key, base_url=base_url)
        self.session_pool = session_pool
        self.api_key = api_key
        self.base_url = base_url

    @retry(
        stop=stop_after_attempt(2),
        wait=wait_exponential(multiplier=1, min=2, max=10),
        retry=retry_if_exception_type((aiohttp.ClientError, asyncio.TimeoutError, Exception)),
        reraise=True
    )
    async def execute_workflow_with_retry(self, inputs: dict, bh: str, timeout: int, error_callback) -> str:
        try:
            return await self.execute_workflow(inputs, bh, timeout, error_callback)
        except Exception as e:
            error_msg = str(e)
            if "rate limiting" in error_msg.lower() or "429" in error_msg:
                await error_callback(bh, f"触发速率限制，等待重试: {error_msg}")
                await asyncio.sleep(2)
            else:
                await error_callback(bh, f"工作流执行失败，准备重试: {error_msg}")
            raise

    async def execute_workflow(self, inputs: dict, bh: str, timeout: int, error_callback) -> str:
        try:
            workflow_start_time = time.time()
            
            url = f"{self.base_url}/workflows/run"
            headers = {
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json",
                "Accept": "text/event-stream"
            }
            
            try:
                json_payload = {
                    "inputs": inputs,
                    "user": "me",  # 恢复到之前您接受的修改，即user为"me"
                    "response_mode": "streaming",
                    "timeout": timeout
                }
                payload_size = len(json.dumps(json_payload).encode('utf-8'))
                print(f"Workflow request payload size: {payload_size / (1024 * 1024):.2f} MB")

                async with self.session_pool.post(
                    url,
                    json=json_payload,
                    headers=headers,
                    timeout=aiohttp.ClientTimeout(total=timeout)
                ) as response:
                    if response.status == 429:
                        error_msg = "请求被速率限制，需要等待"
                        await error_callback(bh, error_msg)
                        raise Exception(error_msg)
                        
                    if response.status != 200:
                        error_msg = f"工作流请求失败: HTTP {response.status}"
                        await error_callback(bh, error_msg)
                        raise Exception(error_msg)
                    
                    final_output = ""
                    buffer = ""
                    
                    async for chunk in response.content.iter_chunked(1024):
                        try:
                            if not chunk:
                                continue
                                
                            chunk_str = chunk.decode('utf-8')
                            buffer += chunk_str
                            
                            while '\n\n' in buffer:
                                message, buffer = buffer.split('\n\n', 1)
                                
                                for line in message.split('\n'):
                                    if line.startswith('data: '):
                                        try:
                                            data = json.loads(line[6:])
                                            event = data.get("event")
                                            
                                            if event == WorkflowEvent.WORKFLOW_FINISHED:
                                                if "outputs" in data.get('data', {}):
                                                    outputs = data['data']['outputs']
                                                    if isinstance(outputs, dict) and 'output' in outputs:
                                                        final_output = outputs['output']
                                                        workflow_processing_time = time.time() - workflow_start_time
                                                        return final_output
                                        except json.JSONDecodeError:
                                            continue
                                            
                        except Exception as e:
                            await error_callback(bh, f"处理响应数据错误: {str(e)}")
                            continue
                    
                    if not final_output:
                        error_msg = "工作流执行失败，未获取到输出"
                        await error_callback(bh, error_msg)
                        raise Exception(error_msg)
                    
                    return final_output
                    
            except aiohttp.ClientError as e:
                error_msg = f"HTTP请求错误: {str(e)}"
                await error_callback(bh, error_msg)
                raise
            except asyncio.TimeoutError as e:
                error_msg = f"请求超时: {str(e)}"
                await error_callback(bh, error_msg)
                raise
                    
        except Exception as e:
            error_msg = f"工作流执行错误: {str(e)}"
            await error_callback(bh, error_msg)
            raise 